Previously we will get gpt-oss model with error layout from examples/converters/convert_megatron_to_hf.py, this PR will fix it. See NVIDIA-NeMo/Megatron-Bridge#3271 for more details.

Validate Steps:

1. Import hf to megatron, train one step, and save the megatron ckpt.

   NRL_FORCE_REBUILD_VENVS=true \
   uv run python examples/run_grpo.py \
       --config examples/configs/recipes/llm/grpo-gptoss-20b-8n8g-megatron.yaml \
       grpo.max_num_steps=1 \
       policy.max_total_sequence_length=512 \
       logger.wandb_enabled=false \
       logger.tensorboard_enabled=false \
       checkpointing.enabled=True \
       checkpointing.checkpoint_dir=results/grpo-gptoss-20b-8n8g-megatron-test-export-transpose \
       checkpointing.save_period=1
   

2. Convert saved megatron ckpt to hf.

   uv run --extra mcore python examples/converters/convert_megatron_to_hf.py \
       --config results/grpo-gptoss-20b-8n8g-megatron-test-export-transpose/step_1/config.yaml \
       --hf-model-name unsloth/gpt-oss-20b-BF16 \
       --megatron-ckpt-path results/grpo-gptoss-20b-8n8g-megatron-test-export-transpose/step_1/policy/weights/iter_0000000 \
       --hf-ckpt-path results/step_1_hf
   

3. Use the converted hf ckpt to train again.

   uv run python examples/run_grpo.py \
       --config examples/configs/recipes/llm/grpo-gptoss-20b-8n8g-megatron.yaml \
       policy.model_name=results/step_1_hf \
       grpo.max_num_steps=1 \
       policy.max_total_sequence_length=512 \
       logger.wandb_enabled=false \
       logger.tensorboard_enabled=false \
       checkpointing.enabled=false
   

Results of "Validate Steps" 3:
Before this PR:

  • Generation KL Error: 13.0520
  • Avg Reward: 0.0000

After this PR:

  • Generation KL Error: 0.0009
  • Avg Reward: 0.3960